/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.berkeley.nlp.util.Pair;
import edu.umd.clip.lm.util.BinaryPrefix;
import edu.umd.clip.lm.util.SimpleList;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.factors.FactorTuple;
import edu.umd.clip.lm.factors.FactorTupleDescription;
import edu.umd.clip.lm.model.*;
import edu.umd.clip.lm.model.decoding.Tagger.TaggerNbestItem;
import edu.umd.clip.lm.model.decoding.VariationalTaggerTrellisSlice.Trie;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class VariationalTagger extends Tagger {
	final int taggerCtxSize;
	
	public static class VariationalPath implements Comparable<VariationalPath>, TaggedPath {
		final double logProb;
		final SimpleList<Integer> tags;
		private static byte tagFactor;
		
		public VariationalPath(double logProb, SimpleList<Integer> tags) {
			this.logProb = logProb;
			this.tags = tags;
		}

		public VariationalPath append(int tag, double logProb) {
			VariationalPath variationalPath = new VariationalPath(this.logProb + logProb, tags.add(tag));
			return variationalPath;
		}
		
		@Override
		public int compareTo(VariationalPath other) {
			return (int) Math.signum(logProb - other.logProb);
		}
		
		public int[] getContext(int ctxSize) {
			int ctx[] = new int[ctxSize];
			SimpleList<Integer> lst = tags;
			for(int i=0; i<ctxSize; ++i) {
				if (lst.head() != null) {
					ctx[i] = lst.head().intValue();
					lst = lst.tail();
				} else {
					ctx[i] = Dictionary.getStart();
				}
			}
			return ctx;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getLogProb()
		 */
		public double getLogProb() {
			return logProb;
		}

		/* (non-Javadoc)
		 * @see edu.umd.clip.lm.model.decoding.TaggedPath#getOrderedTags()
		 */
		public int[] getOrderedTags() {
			int t[] = new int[tags.length()];
			SimpleList<Integer> tail = tags;
			for(int i=1; i<=t.length; ++i) {
				// reverse to the proper order
				t[t.length-i] = tail.head().intValue();
				tail = tail.tail();
			}
			return t;
		}
		
		public String toString() {
			StringBuilder sb = new StringBuilder();
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(tagFactor);
			
			int tags[] = getOrderedTags();
			for(int tag : tags) {
				sb.append(dict.getWord(tag));
				sb.append(' ');
			}
			sb.append('(');
			sb.append(logProb);
			sb.append(')');
			return sb.toString();
		}
	}
	
	public VariationalTagger(Decoder decoder, int beamSize, String tagFactor) {
		super(decoder, beamSize, tagFactor);
		VariationalPath.tagFactor = tagFactorIdx;
		taggerCtxSize = decoder.getForest().getOrder() - 1;
	}

	public Pair<SortedSet<? extends TaggedPath>,Double> evaluateSentence(long[] sentence) {
		ForestModel forest = decoder.getForest();
		TreeTrellisNode currentNode = Decoder.createTreeTrellis(forest.getOrder() - 1);

		final FactorTupleDescription desc = Experiment.getInstance().getTupleDescription();
		final long overtFactorsMask = desc.getOvertFactorsMask();
		final long hiddenFactorsMask = desc.getHiddenFactorsMask();
		
		final int lockNodes = /*forest.getOrder() +*/ 2;
		
		double prob = 1.0;
		TaggerTrellisSlice prevTaggerSlice = new TaggerTrellisSlice();
		{
			// make the start state
			int ctx[] = new int[taggerCtxSize+1];
			Arrays.fill(ctx, Dictionary.getStart());
			prevTaggerSlice.addProb(ctx, 1.0);
			TaggerState state = prevTaggerSlice.getState(ctx);
			state.nbest = new TreeSet<TaggerNbestItem>();
			state.nbest.add(new TaggerNbestItem(null, 0.0, new SimpleList<Integer>()));
		}
		
		TreeSet<TaggerNbestItem> overallNbest = null; 
		for(int i=0; i<sentence.length; ++i) {
			long word = sentence[i] & overtFactorsMask;
			
			if (true || i == sentence.length - 1) {
				overallNbest = new TreeSet<TaggerNbestItem> ();
			}
			
			currentNode = currentNode.appendNodeAndLock(word);
			TreeTrellisNode prevNode = currentNode.getPreviousNode();
			try {
				if (!currentNode.isDecoded()) {
					prevNode.lock(forest.getOrder());
					try {
						DecodingContext context = new DecodingContext(currentNode, forest);
						
						decoder.partition(context);
						
						IdentityHashMap<TrellisNodeItem, Integer> item2tag = decoder.ensurePrefixes(currentNode, tags);
						
						currentNode.finishDecoding();
						
						TaggerTrellisSlice taggerSlice = computeApproximatedProbs(currentNode, item2tag);
						taggerSlice.makeNbest(prevTaggerSlice, beamSize, overallNbest);
						/*
						for(Iterator<TaggerNbestItem> it = overallNbest.descendingIterator(); it.hasNext(); ) {
							TaggerNbestItem item = it.next();
							System.err.println(item.toString());
						}
						*/
						prevTaggerSlice = taggerSlice;
					} finally {
						prevNode.unlock(forest.getOrder());
					}
				}
	
				prob = currentNode.getCumulativeLogScale() + Math.log10(currentNode.getTotalProb());
			} finally {
				currentNode.unlock();
			}
			
		}
		currentNode.releasePath();
		
		return new Pair<SortedSet<? extends TaggedPath>,Double>(overallNbest, prob);
	}
	
	private TaggerTrellisSlice computeApproximatedProbs(TreeTrellisNode currentNode, IdentityHashMap<TrellisNodeItem, Integer> item2tag) {
		TaggerTrellisSlice taggerSlice = new TaggerTrellisSlice();
		int ctx[] = new int[taggerCtxSize+1];
		if (taggerCtxSize > 0) {
			//TreeTrellisNode prevNode = currentNode.getPreviousNode();
			for(TrellisNodeItem item : currentNode.items.keySet()) {
				int tag = getItemTag(item2tag, item);
				if (tag == -1) continue;
				ctx[0] = tag;
				iterate(item, item2tag, taggerSlice, ctx, 1, item.tree.getTotalProb());
			}
		} else {
			// special case for unigram tagger should it ever be needed
			for(TrellisNodeItem item : currentNode.items.keySet()) {
				int tag = getItemTag(item2tag, item);
				if (tag == -1) continue;
				ctx[0] = tag;
				taggerSlice.addProb(ctx, item.tree.getTotalProb());
			}
		}
		return taggerSlice;
	}

	private void iterate(final TrellisNodeItem frontItem, final IdentityHashMap<TrellisNodeItem, Integer> item2tag, 
			final TaggerTrellisSlice taggerSlice, final int ctx[], 
			final int pos, final double prob) 
	{
		for(TrellisNodeItem prevItem : frontItem.getPrevItems()) {
			int tag = getItemTag(item2tag, prevItem);
			if (tag == -1) continue;
			ctx[pos] = tag;
			double newProb = prob * prevItem.tree.getTotalProb();
			if (pos == ctx.length-1) {
				taggerSlice.addProb(ctx, newProb);
			} else {
				iterate(prevItem, item2tag, taggerSlice, ctx, pos+1, newProb);
			}
		}
	}
/*
	private void iterate(final TrellisNodeItem frontItem, final IdentityHashMap<TrellisNodeItem, Integer> item2tag, 
			final VariationalTaggerTrellisSlice taggerSlice, Trie trie, int stepsLeft, final TrellisItem item, double prob) 
	{
		if (stepsLeft == 0) {
			final TrellisSlice slice = frontItem.slice.getNext();
			Collection<TrellisItem> nextItems = slice.prev2curr.get(frontItem);
			for(TrellisItem nextItem : nextItems) {
				int tag = getItemTag(item2tag, nextItem);
				if (tag == -1) continue;
				trie.addProb(tag, prob * nextItem.tree.getTotalProb());
			}
			return;
		}
		
		Collection<TrellisItem> prevs = null;
		
		if (item != null) {
			prevs = item.slice.getPreviousItems(item);
		}
		
		if (prevs == null) {
			int tag = Dictionary.getStart();
			Trie subTrie = trie.addProb(tag, prob);
			iterate(frontItem, context, taggerSlice, subTrie, stepsLeft-1, null, prob);
			return;
		}
		
		for(TrellisItem prevItem : prevs) {
			int tag = context.getItemTag(prevItem);
			if (tag == -1) continue;
			Trie subTrie = trie.addProb(tag, prob * prevItem.tree.getTotalProb());
			iterate(frontItem, context, taggerSlice, subTrie, stepsLeft-1, prevItem, prob * prevItem.tree.getTotalProb());
		}
	}
*/
	protected int getItemTag(IdentityHashMap<TrellisNodeItem, Integer> item2tag, TrellisNodeItem item) {
		Integer itemTag = item2tag.get(item);
		if (itemTag == null) {
			for(int tag=0; tag < tags.length; ++tag) {
				BinaryPrefix tagPrefix = tags[tag];
				if (tagPrefix == null) continue;
				if (tagPrefix.isPrefixOf(item.tree.getPrefix())) {
					itemTag = Integer.valueOf(tag);
					item2tag.put(item, itemTag);
					return tag;
				}
			}
			//System.out.printf("item %s is not a tag: treeProb=%g\n", item, item.tree.getProbTree().getTreeProb());
			return -1;
		}
		return itemTag.intValue();
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.decoding.Tagger#evaluateSentence(java.util.Collection)
	 */
	@Override
	public Pair<SortedSet<? extends TaggedPath>, Double> evaluateSentence(Collection<FactorTuple> sentence) {
		long sent[] = new long[sentence.size()];
		int i=0;
		for(FactorTuple tuple : sentence) {
			sent[i++] = tuple.getBits();
		}
		return evaluateSentence(sent);
	}
}
